"""
Author: Morphlng
Date: 2024-03-19 20:56:58
LastEditTime: 2024-03-25 10:36:41
LastEditors: Morphlng
Description: Test all available sensors in carla
FilePath: /DrivingGym/tests/carla/carla_sensor_test.py
"""

from driving_gym.simulation.adapter.carla import CarlaActorConfig, CarlaConfig, CarlaSensorConfig
from driving_gym.simulation.adapter_factory import AdapterFactory

config = (
    CarlaConfig()
    .connection(host="localhost", port=2000)
    .synchronization(sync_mode=True, fixed_delta_seconds=0.05)
)

adapter = AdapterFactory.create(config.to_dict())


def test_init_adapter():
    adapter.initialize()
    assert adapter._client is not None, "client is not initialized"


def test_reload_world():
    adapter.reload_world()
    assert adapter._map is None, "Adapter data should be reset after reloading world"


def test_create_hero():
    adapter.create_actor(
        "hero",
        CarlaActorConfig().blueprint("vehicle.tesla.model3", rolename="hero").to_dict(),
    )

    assert adapter.get_transform("hero") is not None, "hero is not created"


def test_rgb_sensor():
    camera_config = (
        CarlaSensorConfig()
        .blueprint(type="rgb", attributes={"image_size_x": 800, "image_size_y": 600})
        .spawning(parent="hero", attachment_type="spring_arm")
        .position(relative_x=(-2, 0.5), relative_z=(2, 0.5), relative_pitch=8.0)
    )

    adapter.create_sensor("rgb", camera_config.to_dict())
    for _ in range(10):
        adapter.tick()

    rgb_data = adapter.get_sensor_data("hero:rgb")
    assert rgb_data is not None, "camera sensor is not created"
    assert rgb_data.shape == (600, 800, 3), "camera sensor data shape is not correct"


def test_semantic_sensor():
    camera_config = (
        CarlaSensorConfig()
        .blueprint(type="semseg", attributes={"image_size_x": 512, "image_size_y": 512})
        .spawning(parent="hero")
    )

    adapter.create_sensor("segment", camera_config.to_dict())
    for _ in range(10):
        adapter.tick()

    segment_data = adapter.get_sensor_data("hero:segment")
    assert segment_data is not None, "segment sensor is not created"
    assert segment_data.shape == (
        512,
        512,
        3,
    ), "segment sensor data shape is not correct"


def test_depth_sensor():
    camera_config = (
        CarlaSensorConfig()
        .blueprint(type="depth", attributes={"image_size_x": 512, "image_size_y": 512})
        .spawning(parent="hero")
    )

    adapter.create_sensor("depth", camera_config.to_dict())
    for _ in range(10):
        adapter.tick()

    depth_data = adapter.get_sensor_data("hero:depth")
    assert depth_data is not None, "depth sensor is not created"
    assert depth_data.shape == (512, 512, 3), "depth sensor data shape is not correct"


def test_lidar_sensor():
    lidar_config = (
        CarlaSensorConfig()
        .blueprint(type="lidar", attributes={"range": 50})
        .spawning(parent="hero")
    )

    adapter.create_sensor("lidar", lidar_config.to_dict())
    for _ in range(10):
        adapter.tick()

    lidar_data = adapter.get_sensor_data("hero:lidar")
    assert lidar_data is not None, "lidar sensor is not created"


def test_gnss_sensor():
    gnss_config = CarlaSensorConfig().blueprint(type="gnss").spawning(parent="hero")

    adapter.create_sensor("gnss", gnss_config.to_dict())
    for _ in range(10):
        adapter.tick()

    gnss_data = adapter.get_sensor_data("hero:gnss")
    assert gnss_data is not None, "gnss sensor is not created"


def test_imu_sensor():
    imu_config = CarlaSensorConfig().blueprint(type="imu").spawning(parent="hero")

    adapter.create_sensor("imu", imu_config.to_dict())
    for _ in range(10):
        adapter.tick()

    imu_data = adapter.get_sensor_data("hero:imu")
    assert imu_data is not None, "imu sensor is not created"


def test_radar_sensor():
    radar_config = (
        CarlaSensorConfig()
        .blueprint(type="radar", attributes={"range": 50})
        .spawning(parent="hero")
    )

    adapter.create_sensor("radar", radar_config.to_dict())
    for _ in range(10):
        adapter.tick()

    radar_data = adapter.get_sensor_data("hero:radar")
    assert radar_data is not None, "radar sensor is not created"


def test_collision_sensor():
    collision_config = (
        CarlaSensorConfig().blueprint(type="collision").spawning(parent="hero")
    )

    adapter.create_sensor("collision", collision_config.to_dict())
    for _ in range(10):
        adapter.tick()

    collision_data = adapter.get_sensor_data("hero:collision")
    assert collision_data is not None, "collision sensor is not created"


def test_lane_sensor():
    lane_config = CarlaSensorConfig().blueprint(type="lane_invasion").spawning(parent="hero")

    adapter.create_sensor("lane_invasion", lane_config.to_dict())
    for _ in range(10):
        adapter.tick()

    lane_data = adapter.get_sensor_data("hero:lane_invasion")
    assert lane_data is not None, "lane sensor is not created"


def test_obstacle_sensor():
    obstacle_config = (
        CarlaSensorConfig()
        .blueprint(type="obstacle", attributes={"range": 50})
        .spawning(parent="hero")
    )

    adapter.create_sensor("obstacle", obstacle_config.to_dict())
    for _ in range(10):
        adapter.tick()

    obstacle_data = adapter.get_sensor_data("hero:obstacle")
    assert obstacle_data is not None, "obstacle sensor is not created"


def test_birdeye_sensor():
    birdeye_config = (
        CarlaSensorConfig()
        .blueprint(type="bev", attributes={"range": 50})
        .spawning(parent="hero")
    )

    adapter.create_sensor("birdeye", birdeye_config.to_dict())
    birdeye_data = adapter.get_sensor_data("hero:birdeye")
    assert birdeye_data is not None, "birdeye sensor is not created"


def test_cleanup():
    adapter.cleanup()
    assert adapter._client is None, "Client should be None"
